import pandas as pd


# 寻找频繁项集
def find_items(df, min_support):
    # 单项集支持度筛选
    support_dict = {(k,): v for k, v in (df.sum() / len(df)).to_dict().items() if v > min_support}
    items = list(support_dict.keys())
    # 第几次遍历
    cn = 0

    while len(items) > 1:
        cn += 1
        items = gen_candidate_sets(items)  # 生产k+1项候选集
        print('---------------------------------------')
        print(f'候选集C{cn}:')
        for item in items:
            print(item)
        print(f'共{len(items)}项')

        # 新的支持度函数
        # 计算k+1项集(即对应k+1列)的连乘
        support_dict_tmp = {cols: df[list(cols)].prod(axis=1, numeric_only=True).sum() / len(df)
                            for cols in items}
        # 筛选支持度，获取k+1项频繁项集
        items = [k for k, v in support_dict_tmp.items() if v > min_support]

        # 更新支持度字典
        support_dict.update(support_dict_tmp)

    return support_dict


def transaction_encoder(file):
    with open(file) as fp:
        tmp_items = [{item: True for item in line.strip().split('\t')} for line in fp]
    return pd.DataFrame(tmp_items).fillna(False)


def gen_candidate_sets(x):
    x = list(map(sorted, x))
    r = []
    for i in range(len(x)):
        for j in range(i + 1, len(x)):
            if x[i][:-1] == x[j][:-1] and x[i][-1] != x[j][-1]:
                r.append(tuple(x[i][:-1] + sorted([x[j][-1], x[i][-1]])))
    return r


def find_rule(support_dict, confidence):
    result = {}

    for k, v in support_dict.items():
        if len(k) < 2:
            continue
        # 遍历每一种可能的规则
        for i in range(len(k)):
            cond = k[:i] + k[i + 1:]
            cond_str = ', '.join(cond)
            r = k[i]
            tmp_confidence = support_dict[k] / support_dict[cond]
            if tmp_confidence > confidence:
                result[f'{cond_str}  ===>  {r}'] = dict(confidence=tmp_confidence, support=support_dict[k])

    sorted(result.items(), key=lambda x: (x[1]['confidence'], x[1]['support']), reverse=True)
    return result


def main():
    products = transaction_encoder('./dataset/basket-transaction.txt')
    items = find_items(products, 0.1)
    rules = find_rule(items, 0.5)
    print('\n挖掘结果:')
    for rule, index in rules.items():
        print(f'{rule:^40s} 指标: 置信度={index["confidence"]: .6f}, 支持度={index["support"]: .6f}')


if __name__ == '__main__':
    main()
